import cv2
import numpy as np
import math

mtx = np.array([[2.40364914e+03, 0.00000000e+00, 9.57470815e+02],
                [0.00000000e+00, 2.39630877e+03, 5.59335485e+02],
                [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]])
dist = np.array([[-4.03724657e-01, -1.73243467e+00, -5.25914647e-03, 2.61270024e-03, 9.83021357e+00]])
FACE_HALF_HIGH = 18.9
FACE_HALF_WEIGHT = 18.9
image_point = []
Camera_intrinsic = {"mtx": mtx, "dist": dist, }
cap = cv2.VideoCapture(0)
cap.set(3,1920)
cap.set(4,1080)
while True:
    tet,frame1 = cap.read()
    HSV_img = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    dist,binary = cv2.threshold(HSV_img,100,255,cv2.THRESH_BINARY)
    # cv2.namedWindow('dst', cv2.WINDOW_AUTOSIZE)
    # cv2.imshow('dst', binary)
    dst, binary = cv2.threshold(HSV_img, 100, 255, cv2.THRESH_BINARY)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)  # 开操作

    contours, _ = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    x = 0
    y = 0
    max_w = 0
    max_h = 0
    max = -1
    for cont in contours:
        distance = 0
        image_point = []
        obj_points = []

        # 外接矩形
        x, y, w, h = cv2.boundingRect(cont)
        # 在原图上画出预测的矩形
        area = w * h
        print(area)
        if (area > 25000 and area<200000):
            cv2.rectangle(frame1, (x, y), (x + w, y + h), (0, 0, 255), 5)
            print(w, h)

            t = (x, y)
            image_point.append(t)
            t = (x + w, y)
            image_point.append(t)
            t = (x, y + h)
            image_point.append(t)
            t = (x + w, y + h)
            image_point.append(t)
            image_points = np.array((image_point[-4], image_point[-3], image_point[-2], image_point[-1]), dtype=np.float64)

            o = (-FACE_HALF_WEIGHT, -FACE_HALF_HIGH, 0)
            obj_points.append(o)
            o = (FACE_HALF_WEIGHT, -FACE_HALF_HIGH, 0)
            obj_points.append(o)
            o = (-FACE_HALF_WEIGHT, FACE_HALF_HIGH, 0)
            obj_points.append(o)
            o = (FACE_HALF_WEIGHT, FACE_HALF_HIGH, 0)
            obj_points.append(o)
            obj_point = np.array((obj_points[-4], obj_points[-3], obj_points[-2], obj_points[-1]), dtype=np.float64)
            obj_point = obj_point * 2.9 / 3
        
            _, rvec, tvec = cv2.solvePnP(obj_point, image_points, Camera_intrinsic["mtx"], Camera_intrinsic["dist"])
            print(tvec[0], tvec[1])
            distance = math.sqrt(tvec[0] ** 2 + tvec[1] ** 2 + tvec[2] ** 2)  # 计算距离
            distance = distance + distance / 6
            if distance >= 250 and distance < 300:
                distance += 5
            elif distance >= 300:
                distance += 10
            print(distance)
            cv2.putText(frame1, "%.2fcm" % (distance),(x,y), cv2.FONT_HERSHEY_SIMPLEX,2.0, (0, 255, 0), 3)
            cv2.imshow("this",frame1)
            if cv2.waitKey(100) & 0xff == ord('q'): break            
